package com.w.picture.manager.sharding;

import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import com.w.picture.mode.entity.Space;
import com.w.picture.mode.enums.SpaceLevelEnum;
import com.w.picture.mode.enums.SpaceTypeEnum;
import com.w.picture.service.SpaceService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 动态分片管理
 * 1.将管理器注册为Bean,通过@PostConstruct注解,在Bean加载后获取所有的分表并更新配置
 * 2.编写获取分表列表的方法,从数据库中查询符合要求的空间列表,在补充上逻辑表,就得到了完整的分表列表.
 * 3.更新 ShardingSphere 的 actual-data-nodes 动态表名配置 . 获取到 ShardingSphere 的 ContextManager ，找到配置文件中的那条规则进行更新即可.
 * 4.动态创建分表:在分表管理器中新增动态创建分表的方法,通过拼接SQL的方式创建出和picture表结构一样的分表,创建新的分表后记得更新分表节点.
 *
 * 参考文章:
 * 1.关于动态更新 actual-data-nodes: https://www.yuque.com/linghengqian/meve2v/cgi5en
 * 2.相关issues:https://github.com/apache/shardingsphere/issues/21503
 * 3.开源社区的讨论:https://github.com/apache/shardingsphere/discussions/12258#discussioncomment-3917927
 * -----
 * 4.https://github.com/apache/shardingsphere/issues/16725
 * @author wzp
 * @since 2025/2/1
 */
@Slf4j
//@Component
public class DynamicShardingManager {

    @Resource
    private DataSource dataSource;

    @Resource
    private SpaceService spaceService;

    private static final String LOGIC_TABLE_NAME = "picture";

    private static final String DATABASE_NAME = "logic_db"; // 配置文件中的数据库名称

    @PostConstruct
    public void initialize() {
        log.info("初始化动态分表配置...");
        updateShardingTableNodes();
    }

    /**
     * 创建分片表
     * 注意: 想要使用MyBatis Plus 的SqlRunner执行SQL,必须在配置文件中开启一下配置:
     * mybatis-plus:
     * global-config:
     * enable-sql-runner: true
     *
     * @param space
     */
    public void createSpacePictureTable(Space space) {
        // 动态创建分表
        // 仅为旗舰版团队空间创建分表
        if (space.getSpaceType() == SpaceTypeEnum.TEAM.getValue() && space.getSpaceLevel() == SpaceLevelEnum.FLAGSHIP.getValue()) {
            Long spaceId = space.getId();
            String tableName = "picture_" + spaceId;
            // 创建新表
            String createTableSql = "CREATE TABLE " + tableName + " LIKE picture";
            try {
                SqlRunner.db().update(createTableSql);
                // 更新分表
                updateShardingTableNodes();
            } catch (Exception e) {
                log.error("创建图片空间分表失败，空间 id = {}", space.getId());
            }
        }
    }


    /**
     * 更新ShardingSphere的实际数据节点（actual-data-nodes）配置。
     * 此方法首先获取所有动态分表的名称，然后构建新的实际数据节点字符串，
     * 最后更新ShardingSphere中的配置以反映这些变化。
     */
    private void updateShardingTableNodes() {
        // 获取所有符合要求的分表名集合，包括初始表picture及根据spaceId生成的分表
        Set<String> tableNames = fetchAllPictureTableNames();

        // 构建新的实际数据节点字符串，格式为"w_picture.picture_1,w_picture.picture_2,..."
        String newActualDataNodes = tableNames.stream()
                .map(tableName -> "w_picture." + tableName) // 确保前缀合法，指向正确的数据库和表
                .collect(Collectors.joining(","));
        log.info("动态分表 actual-data-nodes 配置: {}", newActualDataNodes); // 记录日志

        // 获取ShardingSphere的ContextManager实例，用于修改分片规则
        ContextManager contextManager = getContextManager();

        // 通过ContextManager获取当前数据库的规则元数据
        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATABASE_NAME)
                .getRuleMetaData();

        // 查找当前数据库中是否存在ShardingRule
        Optional<ShardingRule> shardingRule = ruleMetaData.findSingleRule(ShardingRule.class);
        if (shardingRule.isPresent()) {
            // 如果存在，获取其配置
            ShardingRuleConfiguration ruleConfig = (ShardingRuleConfiguration) shardingRule.get().getConfiguration();

            // 根据逻辑表名更新实际的数据节点配置
            List<ShardingTableRuleConfiguration> updatedRules = ruleConfig.getTables()
                    .stream()
                    .map(oldTableRule -> {
                        // 对于逻辑表名为"picture"的配置项进行更新
                        if (LOGIC_TABLE_NAME.equals(oldTableRule.getLogicTable())) {
                            // 创建新的表规则配置，并设置新的实际数据节点
                            ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newActualDataNodes);

                            // 复制旧的配置策略到新的配置中
                            newTableRuleConfig.setDatabaseShardingStrategy(oldTableRule.getDatabaseShardingStrategy());
                            newTableRuleConfig.setTableShardingStrategy(oldTableRule.getTableShardingStrategy());
                            newTableRuleConfig.setKeyGenerateStrategy(oldTableRule.getKeyGenerateStrategy());
                            newTableRuleConfig.setAuditStrategy(oldTableRule.getAuditStrategy());

                            return newTableRuleConfig; // 返回新的配置
                        }
                        return oldTableRule; // 如果不是目标逻辑表，则返回原配置
                    })
                    .collect(Collectors.toList());

            // 更新规则配置列表
            ruleConfig.setTables(updatedRules);

            // 应用更新后的规则配置到指定数据库
            contextManager.alterRuleConfiguration(DATABASE_NAME, Collections.singleton(ruleConfig));

            // 重新加载数据库以应用更改
            contextManager.reloadDatabase(DATABASE_NAME);

            log.info("动态分表规则更新成功！"); // 成功日志
        } else {
            log.error("未找到 ShardingSphere 的分片规则配置，动态分表更新失败。"); // 错误日志
        }
    }

    /**
     * 获取所有动态表名，包括初始表 picture 和分表 picture_{spaceId}
     */
    private Set<String> fetchAllPictureTableNames() {
        // 为了测试方便，直接对所有团队空间分表（实际上线改为仅对旗舰版生效）
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list()
                .stream()
                .map(Space::getId)
                .collect(Collectors.toSet());
        Set<String> tableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId)
                .collect(Collectors.toSet());
        tableNames.add(LOGIC_TABLE_NAME); // 添加初始逻辑表
        return tableNames;
    }

    /**
     * 获取 ShardingSphere ContextManager
     */
    private ContextManager getContextManager() {
        try (ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)) {
            return connection.getContextManager();
        } catch (SQLException e) {
            throw new RuntimeException("获取 ShardingSphere ContextManager 失败", e);
        }
    }

}

